package a4;

import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;

public class WordTagContainer {
	SortedMap<String, SortedMap<String, Integer> > container;
	SortedMap<String, String> baselineWordTags;
	SortedMap<String, Integer> totalTagCounts;
	boolean containerChanged = false;
	
	public WordTagContainer(){
		container = new TreeMap<String,SortedMap<String, Integer> >();
		baselineWordTags = new TreeMap<String, String>();
		totalTagCounts = new TreeMap<String, Integer>();
	}
	
	public void put(String word, String tag){
		if(container.containsKey(word)){
			SortedMap<String, Integer> tagCount = container.get(word);
			int count = tagCount.containsKey(tag) ? tagCount.get(tag) : 0;
			tagCount.put(tag, count + 1);
			container.put(word, tagCount);
		}else{
			SortedMap<String, Integer> tagCount = new TreeMap<String, Integer>();
			tagCount.put(tag, 1);
			container.put(word, tagCount);
		}
		
		// Increment the total count (i.e. in the entire text) of the given tag
		int count = totalTagCounts.containsKey(tag) ? totalTagCounts.get(tag) : 0;
		totalTagCounts.put(tag, count + 1);
		
		containerChanged = true;
	}
	
	private void updateBaselineTags(){
		Set<String> words = container.keySet();
		for(String w : words){
			SortedMap<String, Integer> countPerTag = container.get(w);
			Set<String> tags = countPerTag.keySet();
			int maxCount = 0;
			String maxCountTag = "";
			for(String t : tags){
				if(countPerTag.get(t) > maxCount){
					maxCount = countPerTag.get(t);
					maxCountTag = t;
				}
			}
			baselineWordTags.put(w, maxCountTag);
		}
		
		// Find the most common tag and set it to the baseline tag of unknown words <unk>
//		Set<String> tags = totalTagCounts.keySet();
//		int maxCount = 0;
//		String maxCountTag = "";
//		for(String t : tags){
//			if(totalTagCounts.get(t) > maxCount){
//				maxCount = totalTagCounts.get(t);
//				maxCountTag = t;
//			}
//		}
//		System.out.println("MAXCOUNTTAG: " + maxCountTag);
//		baselineWordTags.put("<unk>", maxCountTag);
		baselineWordTags.put("<unk>", "NN");
		
		containerChanged = false;
	}
	
	public String getBaselineTag(String word){
		if(containerChanged) updateBaselineTags();
		if(baselineWordTags.containsKey(word)){
			return baselineWordTags.get(word);
		}else{
			return baselineWordTags.get("<unk>");
		}
	}
	
	public String toString(){
		Set<String> words = container.keySet();
		String ret = "";
		for(String w : words){
			SortedMap<String, Integer> tagCount = container.get(w);
			if(tagCount != null){
				Set<String> tags = tagCount.keySet();
				for(String t : tags){
					ret += w + " " + t + " : " + tagCount.get(t) + "\n";
				}

			}
		}
		return ret;
	}

}
